package top.yudoge.dogeeditor.llm.rewrite

import com.google.gson.Gson
import dev.langchain4j.data.message.ChatMessage
import dev.langchain4j.data.message.SystemMessage
import dev.langchain4j.data.message.UserMessage
import dev.langchain4j.model.chat.response.ChatResponse
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.stereotype.Component
import top.yudoge.dogeeditor.commons.utils.toJsonStr
import top.yudoge.dogeeditor.db.api.DataSourceMetaDataApi
import top.yudoge.dogeeditor.llm.memory.SessionMemory
import top.yudoge.dogeeditor.llm.objects.mro.TableMRO
import java.util.LinkedList

@Component
@Deprecated(message = "This class is deprecated, use the new implementation SQLGenerate2ChainNode")
class SQLGenerateChainNode : TaskChainNode {

    @Autowired
    private lateinit var dataSourceMetaDataApi: DataSourceMetaDataApi

    @Autowired
    private lateinit var sessionMemory: SessionMemory

    override fun process(ctx: ChainContext) {
        doProcess(ctx)
    }

    fun doProcess(ctx: ChainContext, lastErrorMessage: String? = null) {
        if (lastErrorMessage!= null) {
            sendReasoningCallback(ctx, "Re-generating SQL...")
        } else {
            sendReasoningCallback(ctx, "Generating SQL...")
        }

        val messages = buildMessages(ctx, lastErrorMessage)
        val streamingModel = ctx.streamingModel
        val callback = ctx.callback

        streamingModel.chat(messages, object : StreamingChatResponseHandler {
            override fun onPartialResponse(p0: String?) {
                if (p0 != null) {
                    callback.token(p0)
                }
            }

            override fun onCompleteResponse(p0: ChatResponse?) {
                p0?.let {
                    val aiMessage = it.aiMessage().text()
                    addChatHistory(ctx, Message("ai", aiMessage))

                    val SQL = extractSQLStatement(aiMessage);
                    if (SQL != null) {
                        putChainStore(ctx, GENERATED_SQL, SQL)
                    }
                }
                if (ctx.query.selfChecking) {
                    emitNext(ctx)
                } else {
                    callback.complete()
                }
            }

            override fun onError(p0: Throwable?) {
                callback.error(p0)
            }
        })

    }

    private fun extractSQLStatement(aiMessage: String): String? {
        val lines = mutableListOf<String>()
        var startFounded = false
        for (line in aiMessage.lines()) {
            if (line.lowercase() == "```sql" && !startFounded) {
                startFounded = true
                continue
            }
            if (line == "```") {
                break
            }
            if (startFounded) {
                lines.add(line)
            }
        }
        return lines.joinToString("\n")
    }


    private fun buildMessages(ctx: ChainContext, lastErrorMessage: String? = null): List<ChatMessage> {

        val userCode = ctx.query.code
        val userTask = ctx.query.userTask
        val dataSource = dataSourceMetaDataApi.getDataSources().first { it.identifier == ctx.query.dataSourceIdentifier }!!

        val relatedTableInfo = getChainStore(ctx, TableInformationGetChainNode.TABLE_INFO) as? List<TableMRO>
        val relatedDocuments = getChainStore(ctx, RAGChainNode.RELATED_DOCS) as? List<String>

        return buildList {
            add(SystemMessage.from("You are a SQL expert and your task is to help users rewrite their SQL."))

            var historyMessageStr: String? = null
            val historyMessage = listChatHistory(ctx)
            if (!historyMessage.isNullOrEmpty()) {
                historyMessageStr = StringBuilder("Here is the recent conversation history(ai is you, user is the user): \n").append(
                    listChatHistory(ctx).toJsonStr()
                ).append("\n\n").toString()
            }

            if (!userCode.isNullOrBlank()) {
                add(UserMessage.from("The user SQL:\n" +
                        "```sql\n" +
                        "${userCode}\n" +
                        "```\n")
                )
            }

            val sb = StringBuilder()

            sb.append("User task: \n" +
                    "```text\n" +
                    "${userTask}\n" +
                    "```\n"
            )

            if (!relatedTableInfo.isNullOrEmpty()) {
                sb.append("Related tables: \n")
                relatedTableInfo.forEach {
                    sb.append(it.toJsonStr())
                    sb.append("\n")
                }
            }

            if (!relatedDocuments.isNullOrEmpty()) {
                sb.append("Related information: \n")
                relatedDocuments.forEach {
                    sb.append(it)
                    sb.append("\n")
                }
            }

            // todo use tool / RAG to support SQL dialect(some new database's documetation may not in the model kownledge)
            sb.append("The database we are using is ${dataSource.vendor.vendorName}, so please use the correct database dialect!\n\n")

            if (lastErrorMessage != null) {
                sb.append("It seems that the last SQL you generated is not correct, so we need to re-generate it. Error message: ${lastErrorMessage}\n\n")
            }
            sb.append("Let's think step by step, give me your thinking process and the final SQL-script. Please return the response in a language that is strictly consistent with the userTask")

            addChatHistory(ctx, Message("user", sb.toString()))

            val newUserMessage = UserMessage.from(
                "${historyMessageStr?:""}${sb.toString()}"
            )
            add(newUserMessage)
        }
    }

    private fun addChatHistory(ctx: ChainContext, message: Message) {
        var historyHolder = getHistoryHolder(ctx)
        historyHolder.add(message)
        sessionMemory.put(ctx.query.sessionId, HISTORY_HOLDER_KEY, Gson().toJson(historyHolder))
    }

    private fun listChatHistory(ctx: ChainContext): List<Message> {
        return getHistoryHolder(ctx).history()
    }

    private fun getHistoryHolder(ctx: ChainContext): HistoryMessageHolder {
        return sessionMemory.get(ctx.query.sessionId, HISTORY_HOLDER_KEY)?.let {
            Gson().fromJson(it, HistoryMessageHolder::class.java)
        } ?: HistoryMessageHolder()
    }

    companion object {
        private const val HISTORY_HOLDER_KEY = "message:history"
        const val GENERATED_SQL = "generated_sql"
    }


    class HistoryMessageHolder {

        private val max = 10
        private val history = LinkedList<Message>()

        fun add(message: Message) {
            if (history.size >= max) {
                history.removeFirst()
            }

            history.addLast(message)
        }

        fun history(): List<Message> {
            return history.toList()
        }
    }

    class Message (
        val type: String,
        val content: String
    )
}